import cv2
import time
import math
from ultralytics import YOLO

class YOLODetector:
    def __init__(self, model_path, confidence_thresh=0.5):
        self.model = YOLO(model_path)
        self.conf_thresh = confidence_thresh

    def predict_image(self, image_path, flight_height, sensor_width, sensor_height, focal_length, pitch_angle=0,
                      roll_angle=0):
        """处理单张图片"""
        results = self.model(image_path, conf=self.conf_thresh)
        return self._parse_results(results, flight_height, sensor_width, sensor_height, focal_length, pitch_angle,
                                   roll_angle)

    def predict_video(self, video_path, flight_height, sensor_width, sensor_height, focal_length, pitch_angle=0,
                      roll_angle=0):
        """处理视频文件"""
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            return {
                "error": "Could not open video file",
                "fps": 0,
                "frame_count": 0,
                "processing_time": 0
            }

        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        video_results = {
            "fps": fps,
            "frame_count": frame_count,
            "processing_time": 0
        }

        start_time = time.time()
        frame_id = 0

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            results = self.model(frame, conf=self.conf_thresh)
            frame_result = self._parse_results(results, flight_height, sensor_width, sensor_height, focal_length,
                                               pitch_angle, roll_angle)
            frame_result["frame_id"] = frame_id

            simplified_result = {
                "frame_id": frame_id,
                "detection_count": len(frame_result["detections"]),
                "gsd_x": frame_result["gsd_x"],
                "gsd_y": frame_result["gsd_y"]
            }

            if len(frame_result["detections"]) > 0:
                simplified_result["sample_detections"] = frame_result["detections"][:5]

            video_results.setdefault("frames", []).append(simplified_result)
            frame_id += 1

        cap.release()
        video_results["processing_time"] = time.time() - start_time
        return video_results

    def _parse_results(self, results, flight_height, sensor_width, sensor_height, focal_length, pitch_angle=0,
                       roll_angle=0):
        """解析YOLO结果并考虑姿态角"""
        # 处理空结果
        if not results or len(results) == 0 or not results[0].boxes:
            return {
                "detections": [],
                "original_size": (0, 0),
                "gsd_x": 0,
                "gsd_y": 0
            }

        image_width, image_height = results[0].orig_shape[1], results[0].orig_shape[0]

        # 毫米转米
        sensor_width_m = sensor_width / 1000.0
        sensor_height_m = sensor_height / 1000.0
        focal_length_m = focal_length / 1000.0

        # 计算有效高度（考虑俯仰角和横滚角）
        effective_height = self.calculate_effective_height(flight_height, pitch_angle, roll_angle)

        # 计算GSD（考虑姿态角）
        gsd_x = (effective_height * sensor_width_m) / (focal_length_m * image_width)
        gsd_y = (effective_height * sensor_height_m) / (focal_length_m * image_height)

        detections = []
        for result in results:
            if not result.boxes:
                continue

            for box in result.boxes:
                xmin, ymin, xmax, ymax = box.xyxy[0].cpu().numpy()

                # 计算物体中心点坐标（归一化到0-1）
                center_x = (xmin + xmax) / (2 * image_width)
                center_y = (ymin + ymax) / (2 * image_height)

                # 计算位置校正因子
                position_correction = self.calculate_position_correction(
                    center_x, center_y,
                    pitch_angle, roll_angle
                )

                # 应用位置校正
                length = (xmax - xmin) * gsd_x * position_correction
                width = (ymax - ymin) * gsd_y * position_correction

                detections.append({
                    "class_id": int(box.cls.cpu().numpy()[0]),
                    "class_name": result.names[int(box.cls.cpu().numpy()[0])],
                    "confidence": float(box.conf.cpu().numpy()[0]),
                    "bbox": {
                        "xmin": float(xmin),
                        "ymin": float(ymin),
                        "xmax": float(xmax),
                        "ymax": float(ymax)
                    },
                    "length": float(length),
                    "width": float(width),
                    "position_correction": float(position_correction)
                })

        return {
            "detections": detections,
            "original_size": results[0].orig_shape,
            "gsd_x": float(gsd_x),
            "gsd_y": float(gsd_y),
            "effective_height": float(effective_height)
        }

    def calculate_effective_height(self, height, pitch, roll):
        """计算考虑姿态角后的有效高度"""
        # 将角度转换为弧度
        pitch_rad = math.radians(pitch)
        roll_rad = math.radians(roll)

        # 计算有效高度（考虑两个轴的倾斜）
        return height * math.cos(pitch_rad) * math.cos(roll_rad)

    def calculate_position_correction(self, norm_x, norm_y, pitch, roll):
        """计算物体在图像中位置引起的尺寸校正因子"""
        # 将角度转换为弧度
        pitch_rad = math.radians(pitch)
        roll_rad = math.radians(roll)

        # 计算物体偏离中心的距离（归一化到-0.5到0.5）
        dx = norm_x - 0.5
        dy = norm_y - 0.5

        # 计算位置引起的视角变化
        # 简化模型：假设相机倾斜导致物体位置变化引起尺寸变化
        # 实际应用中可能需要更复杂的几何模型
        return 1.0 / (1.0 + math.tan(pitch_rad) * dy + math.tan(roll_rad) * dx)

    def calculate_corrected_coordinates(self, center_x, center_y, image_width, image_height, gsd_x, gsd_y, longitude,
                                        latitude):
        """计算修正后的经纬度坐标"""
        # 计算图像中心到物体的偏移量（米）
        offset_x = (center_x - 0.5) * image_width * gsd_x
        offset_y = (center_y - 0.5) * image_height * gsd_y

        # 地球半径（米）
        earth_radius = 6371000

        # 计算经纬度偏移量（度）
        # 注意：这里简化处理，假设在小范围内，1度纬度约等于111,111米
        # 经度的缩放因子随纬度变化，cos(latitude)
        lat_per_meter = 1 / 111111.0
        lon_per_meter = 1 / (111111.0 * math.cos(math.radians(latitude)))

        # 计算修正后的经纬度
        corrected_lat = latitude + (offset_y / earth_radius) * (180 / math.pi)
        corrected_lon = longitude + (offset_x / earth_radius) * (180 / math.pi) / math.cos(math.radians(latitude))

        return {
            "longitude": float(corrected_lon),
            "latitude": float(corrected_lat)
        }